package com.boost.tjvm.agent;


import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Modifier;
import java.security.ProtectionDomain;

/**
 * 1. 打包tjvm.jar Tasks->build->jar
 * 2. copy tools.jar javassist-3.28.0-GA.jar 到 Class-Path 目录下
 * 2. 编辑 MANIFEST.MF
 Manifest-Version: 1.0
 Premain-Class: com.boost.tjvm.agent.MyAgent
 Agent-Class: com.boost.tjvm.agent.MyAgent
 Class-Path: tools.jar javassist-3.28.0-GA.jar
 Can-Redefine-Classes: true
 Can-Retransform-Classes: true
 * 4. idea 启动Tast
 * 5. 执行
 * java -cp tjvm.jar com.boost.tjvm.agent.AttachTest 7764 D:\\tjvm.jar
 */
public class MyAgent {
    //使用 Attach API 远程加载的 Java agent 不会再先于main方法执行，这取决于另一虚拟机调用 Attach API 的时机。并且，它运行的也不再是premain方法，而是名为agentmain的方法。
    public static void agentmain(String args, Instrumentation inst) throws Exception {
        //Can-Redefine-Classes 虚拟机没加载过字节码，在首次加载的时候重定义字节码 对应方法premain
        //Can-Retransform-Classes 虚拟机已经加载了字节码，需要改变字节码时候出发对应agentmain
        System.out.println("agentmain");
        Transformer transformer = new Transformer();
        inst.addTransformer(transformer, true);
        inst.retransformClasses(Class.forName(args)); // 允许修改Task类
        Thread.sleep(20000);
        //恢复字节码 对应Can-Retransform-Classes， 此处移除 transformer 所以可以恢复
        inst.removeTransformer(transformer);
        inst.retransformClasses(Class.forName(args)); // 允许修改Task类
    }

    private static class Transformer implements ClassFileTransformer {
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
            className = className.replaceAll("/", ".");
            System.out.println("transform: " + className);

            ClassPool classPool = ClassPool.getDefault();
            classPool.appendClassPath(new LoaderClassPath(loader)); // 将要修改的类的classpath加入到ClassPool中，否则找不到该类
            try {
                CtClass ctClass = classPool.get(className);
                for (CtMethod ctMethod : ctClass.getDeclaredMethods()) {
                    //非main方法
                    if (Modifier.isPublic(ctMethod.getModifiers()) && !ctMethod.getName().equals("main")) {
                        // 修改字节码
                        ctMethod.addLocalVariable("begin", CtClass.longType);
                        ctMethod.addLocalVariable("end", CtClass.longType);
                        ctMethod.insertBefore("begin = System.currentTimeMillis();");
                        ctMethod.insertAfter("end = System.currentTimeMillis();");
                        ctMethod.insertAfter("System.out.println(\"方法" + ctMethod.getName() + "耗时\"+ (end - begin) +\"ms\");");
                    }
                }
                ctClass.detach();
                return ctClass.toBytecode();
            } catch (Exception e) {
                e.printStackTrace();
            }
            return classfileBuffer;
        }
    }

}
